numpy>=1.16
tensorflow==2.1.0
tensorflow_datasets==1.3.2
scikit-learn>=0.20
jax>=0.1.56
-e git://github.com/google-research/flax.git@prerelease#egg=flax
